/**
 * Copyright (c) Huawei Technologies Co., Ltd. 2022. All rights reserved.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
#include "config/openmp.h"
#include "core/mq_base_types.h"
#include "core/utils.h"
#include "math/pr/parameter_resolver.h"
#include "simulator/utils.h"
#ifdef __x86_64__
#    include "simulator/densitymatrix/detail/cpu_densitymatrix_avx_double_policy.h"
#    include "simulator/densitymatrix/detail/cpu_densitymatrix_avx_float_policy.h"
#elif defined(__amd64)
#    include "simulator/densitymatrix/detail/cpu_densitymatrix_arm_double_policy.h"
#    include "simulator/densitymatrix/detail/cpu_densitymatrix_arm_float_policy.h"
#endif
#include "simulator/densitymatrix/detail/cpu_densitymatrix_policy.h"

namespace mindquantum::sim::densitymatrix::detail {
// Single qubit operator
// ========================================================================================================

// method is based on 'mq_vector' simulator, extended to densitymatrix
template <typename derived_, typename calc_type_>
void CPUDensityMatrixPolicyBase<derived_, calc_type_>::ApplySingleQubitMatrix(const qs_data_p_t& src_out,
                                                                              qs_data_p_t* des_p, qbit_t obj_qubit,
                                                                              const qbits_t& ctrls, const matrix_t& m,
                                                                              index_t dim) {
    auto& des = (*des_p);
    if (des == nullptr) {
        des = derived::InitState(dim);
    }
    qs_data_p_t src;
    bool will_free = false;
    if (src_out == nullptr) {
        src = derived::InitState(dim);
        will_free = true;
    } else {
        src = src_out;
    }
    SingleQubitGateMask mask({obj_qubit}, ctrls);
    if (!mask.ctrl_mask) {
        THRESHOLD_OMP_FOR(
            dim, DimTh, for (omp::idx_t a = 0; a < static_cast<omp::idx_t>(dim / 2); a++) {  // loop on the row
                auto r0 = ((a & mask.obj_high_mask) << 1) + (a & mask.obj_low_mask);
                auto r1 = r0 + mask.obj_mask;
                for (index_t b = 0; b <= a; b++) {  // loop on the column
                    auto c0 = ((b & mask.obj_high_mask) << 1) + (b & mask.obj_low_mask);
                    auto c1 = c0 + mask.obj_mask;
                    qs_data_t src_00 = src[IdxMap(r0, c0)];
                    qs_data_t src_11 = src[IdxMap(r1, c1)];
                    qs_data_t src_10 = src[IdxMap(r1, c0)];
                    // for matrix[row, col], only in this case (row < col) is possible
                    qs_data_t src_01 = GetValue(src, r0, c1);
                    auto des_00 = m[0][0] * std::conj(m[0][0]) * src_00 + m[0][0] * std::conj(m[0][1]) * src_01
                                  + m[0][1] * std::conj(m[0][0]) * src_10 + m[0][1] * std::conj(m[0][1]) * src_11;
                    auto des_11 = m[1][0] * std::conj(m[1][0]) * src_00 + m[1][0] * std::conj(m[1][1]) * src_01
                                  + m[1][1] * std::conj(m[1][0]) * src_10 + m[1][1] * std::conj(m[1][1]) * src_11;
                    auto des_01 = m[0][0] * std::conj(m[1][0]) * src_00 + m[0][0] * std::conj(m[1][1]) * src_01
                                  + m[0][1] * std::conj(m[1][0]) * src_10 + m[0][1] * std::conj(m[1][1]) * src_11;
                    auto des_10 = m[1][0] * std::conj(m[0][0]) * src_00 + m[1][0] * std::conj(m[0][1]) * src_01
                                  + m[1][1] * std::conj(m[0][0]) * src_10 + m[1][1] * std::conj(m[0][1]) * src_11;

                    des[IdxMap(r0, c0)] = des_00;
                    des[IdxMap(r1, c1)] = des_11;
                    des[IdxMap(r1, c0)] = des_10;
                    SetValue(des, r0, c1, des_01);
                }
            })
    } else {
        THRESHOLD_OMP_FOR(
            dim, DimTh, for (omp::idx_t a = 0; a < static_cast<omp::idx_t>(dim / 2); a++) {  // loop on the row
                auto r0 = ((a & mask.obj_high_mask) << 1) + (a & mask.obj_low_mask);
                auto r1 = r0 + mask.obj_mask;
                for (index_t b = 0; b <= a; b++) {  // loop on the column
                    auto c0 = ((b & mask.obj_high_mask) << 1) + (b & mask.obj_low_mask);
                    if (((r0 & mask.ctrl_mask) != mask.ctrl_mask)
                        && ((c0 & mask.ctrl_mask) != mask.ctrl_mask)) {  // both not in control
                        continue;
                    }
                    auto c1 = c0 + mask.obj_mask;
                    qs_data_t src_00 = src[IdxMap(r0, c0)];
                    qs_data_t src_11 = src[IdxMap(r1, c1)];
                    qs_data_t src_10 = src[IdxMap(r1, c0)];
                    qs_data_t src_01 = GetValue(src, r0, c1);
                    qs_data_t des_00;
                    qs_data_t des_11;
                    qs_data_t des_01;
                    qs_data_t des_10;
                    if ((r0 & mask.ctrl_mask) == mask.ctrl_mask) {
                        if ((c0 & mask.ctrl_mask) == mask.ctrl_mask) {  // both in control
                            des_00 = m[0][0] * std::conj(m[0][0]) * src_00 + m[0][0] * std::conj(m[0][1]) * src_01
                                     + m[0][1] * std::conj(m[0][0]) * src_10 + m[0][1] * std::conj(m[0][1]) * src_11;
                            des_11 = m[1][0] * std::conj(m[1][0]) * src_00 + m[1][0] * std::conj(m[1][1]) * src_01
                                     + m[1][1] * std::conj(m[1][0]) * src_10 + m[1][1] * std::conj(m[1][1]) * src_11;
                            des_01 = m[0][0] * std::conj(m[1][0]) * src_00 + m[0][0] * std::conj(m[1][1]) * src_01
                                     + m[0][1] * std::conj(m[1][0]) * src_10 + m[0][1] * std::conj(m[1][1]) * src_11;
                            des_10 = m[1][0] * std::conj(m[0][0]) * src_00 + m[1][0] * std::conj(m[0][1]) * src_01
                                     + m[1][1] * std::conj(m[0][0]) * src_10 + m[1][1] * std::conj(m[0][1]) * src_11;
                        } else {  // row in control but not column
                            des_00 = m[0][0] * src_00 + m[0][1] * src_10;
                            des_10 = m[1][0] * src_00 + m[1][1] * src_10;
                            des_01 = m[0][0] * src_01 + m[0][1] * src_11;
                            des_11 = m[1][0] * src_01 + m[1][1] * src_11;
                        }
                    } else {  // column in control but not row
                        des_00 = std::conj(m[0][0]) * src_00 + std::conj(m[0][1]) * src_01;
                        des_01 = std::conj(m[1][0]) * src_00 + std::conj(m[1][1]) * src_01;
                        des_10 = std::conj(m[0][0]) * src_10 + std::conj(m[0][1]) * src_11;
                        des_11 = std::conj(m[1][0]) * src_10 + std::conj(m[1][1]) * src_11;
                    }
                    des[IdxMap(r0, c0)] = des_00;
                    des[IdxMap(r1, c1)] = des_11;
                    des[IdxMap(r1, c0)] = des_10;
                    SetValue(des, r0, c1, des_01);
                }
            })
    }
    if (will_free) {
        derived::FreeState(&src);
    }
}

template <typename derived_, typename calc_type_>
void CPUDensityMatrixPolicyBase<derived_, calc_type_>::ApplyTwoQubitsMatrix(const qs_data_p_t& src, qs_data_p_t* des_p,
                                                                            const qbits_t& objs, const qbits_t& ctrls,
                                                                            const matrix_t& m, index_t dim) {
    DoubleQubitGateMask mask(objs, ctrls);
    if (!mask.ctrl_mask) {
        ApplyTwoQubitsMatrixNoCtrl(src, des_p, objs, ctrls, m, dim);
    } else {
        ApplyTwoQubitsMatrixCtrl(src, des_p, objs, ctrls, m, dim);
    }
}

template <typename derived_, typename calc_type_>
void CPUDensityMatrixPolicyBase<derived_, calc_type_>::ApplyTwoQubitsMatrixNoCtrl(const qs_data_p_t& src_out,
                                                                                  qs_data_p_t* des_p,
                                                                                  const qbits_t& objs,
                                                                                  const qbits_t& ctrls,
                                                                                  const matrix_t& m, index_t dim) {
    auto& des = (*des_p);
    if (des == nullptr) {
        des = derived::InitState(dim);
    }
    qs_data_p_t src;
    bool will_free = false;
    if (src_out == nullptr) {
        src = derived::InitState(dim);
        will_free = true;
    } else {
        src = src_out;
    }
    DoubleQubitGateMask mask(objs, ctrls);
    size_t mask1 = (static_cast<uint64_t>(1) << objs[0]);
    size_t mask2 = (static_cast<uint64_t>(1) << objs[1]);
    THRESHOLD_OMP_FOR(
        dim, DimTh, for (omp::idx_t a = 0; a < static_cast<omp::idx_t>(dim / 4); a++) {
            VT<index_t> row(4);  // row index of reduced matrix entry
            SHIFT_BIT_TWO(mask.obj_low_mask, mask.obj_rev_low_mask, mask.obj_high_mask, mask.obj_rev_high_mask, a,
                          row[0]);
            row[3] = row[0] + mask.obj_mask;
            row[1] = row[0] + mask1;
            row[2] = row[0] + mask2;
            for (index_t b = 0; b <= a; b++) {
                VT<index_t> col(4);  // column index of reduced matrix entry
                SHIFT_BIT_TWO(mask.obj_low_mask, mask.obj_rev_low_mask, mask.obj_high_mask, mask.obj_rev_high_mask, b,
                              col[0]);
                col[3] = col[0] + mask.obj_mask;
                col[1] = col[0] + mask1;
                col[2] = col[0] + mask2;
                VT<VT<qs_data_t>> tmp_mat(4, VT<qs_data_t>(4));
                for (int i = 0; i < 4; i++) {
                    for (int j = 0; j < 4; j++) {
                        tmp_mat[i][j] = m[i][0] * GetValue(src, row[0], col[j])
                                        + m[i][1] * GetValue(src, row[1], col[j])
                                        + m[i][2] * GetValue(src, row[2], col[j])
                                        + m[i][3] * GetValue(src, row[3], col[j]);
                    }
                }

                for (int i = 0; i < 4; i++) {
                    for (int j = 0; j < 4; j++) {
                        auto new_value = tmp_mat[i][0] * std::conj(m[j][0]) + tmp_mat[i][1] * std::conj(m[j][1])
                                         + tmp_mat[i][2] * std::conj(m[j][2]) + tmp_mat[i][3] * std::conj(m[j][3]);
                        SetValue(des, row[i], col[j], new_value);
                    }
                }
            }
        })
    if (will_free) {
        derived::FreeState(&src);
    }
}

template <typename derived_, typename calc_type_>
void CPUDensityMatrixPolicyBase<derived_, calc_type_>::ApplyTwoQubitsMatrixCtrl(const qs_data_p_t& src_out,
                                                                                qs_data_p_t* des_p, const qbits_t& objs,
                                                                                const qbits_t& ctrls, const matrix_t& m,
                                                                                index_t dim) {
    auto& des = (*des_p);
    if (des == nullptr) {
        des = derived::InitState(dim);
    }
    qs_data_p_t src;
    bool will_free = false;
    if (src_out == nullptr) {
        src = derived::InitState(dim);
        will_free = true;
    } else {
        src = src_out;
    }
    DoubleQubitGateMask mask(objs, ctrls);
    size_t mask1 = (static_cast<uint64_t>(1) << objs[0]);
    size_t mask2 = (static_cast<uint64_t>(1) << objs[1]);
    THRESHOLD_OMP_FOR(
        dim, DimTh, for (omp::idx_t a = 0; a < static_cast<omp::idx_t>(dim / 4); a++) {
            VT<index_t> row(4);  // row index of reduced matrix entry
            SHIFT_BIT_TWO(mask.obj_low_mask, mask.obj_rev_low_mask, mask.obj_high_mask, mask.obj_rev_high_mask, a,
                          row[0]);
            row[3] = row[0] + mask.obj_mask;
            row[1] = row[0] + mask1;
            row[2] = row[0] + mask2;
            for (index_t b = 0; b <= a; b++) {
                VT<index_t> col(4);  // column index of reduced matrix entry
                SHIFT_BIT_TWO(mask.obj_low_mask, mask.obj_rev_low_mask, mask.obj_high_mask, mask.obj_rev_high_mask, b,
                              col[0]);
                if (((row[0] & mask.ctrl_mask) != mask.ctrl_mask)
                    && ((col[0] & mask.ctrl_mask) != mask.ctrl_mask)) {  // both not in control
                    continue;
                }
                col[3] = col[0] + mask.obj_mask;
                col[1] = col[0] + mask1;
                col[2] = col[0] + mask2;
                VT<VT<qs_data_t>> tmp_mat(4, VT<qs_data_t>(4));
                if ((row[0] & mask.ctrl_mask) == mask.ctrl_mask) {  // row in control
                    for (int i = 0; i < 4; i++) {
                        for (int j = 0; j < 4; j++) {
                            tmp_mat[i][j] = m[i][0] * GetValue(src, row[0], col[j])
                                            + m[i][1] * GetValue(src, row[1], col[j])
                                            + m[i][2] * GetValue(src, row[2], col[j])
                                            + m[i][3] * GetValue(src, row[3], col[j]);
                        }
                    }
                } else {  // row not in control
                    for (int i = 0; i < 4; i++) {
                        for (int j = 0; j < 4; j++) {
                            tmp_mat[i][j] = GetValue(src, row[i], col[j]);
                        }
                    }
                }
                if ((col[0] & mask.ctrl_mask) == mask.ctrl_mask) {  // column in control
                    for (int i = 0; i < 4; i++) {
                        for (int j = 0; j < 4; j++) {
                            auto new_value = tmp_mat[i][0] * std::conj(m[j][0]) + tmp_mat[i][1] * std::conj(m[j][1])
                                             + tmp_mat[i][2] * std::conj(m[j][2]) + tmp_mat[i][3] * std::conj(m[j][3]);
                            SetValue(des, row[i], col[j], new_value);
                        }
                    }
                } else {  // column not in control
                    for (int i = 0; i < 4; i++) {
                        for (int j = 0; j < 4; j++) {
                            SetValue(des, row[i], col[j], tmp_mat[i][j]);
                        }
                    }
                }
            }
        })
    if (will_free) {
        derived::FreeState(&src);
    }
}

template <typename derived_, typename calc_type_>
void CPUDensityMatrixPolicyBase<derived_, calc_type_>::ApplyNQubitsMatrix(const qs_data_p_t& src_out,
                                                                          qs_data_p_t* des_p, const qbits_t& objs,
                                                                          const qbits_t& ctrls, const matrix_t& gate,
                                                                          index_t dim) {
    auto& des = *des_p;
    if (des == nullptr) {
        des = derived::InitState(dim);
    }
    qs_data_p_t src;
    bool will_free = false;
    if (src_out == nullptr) {
        src = derived::InitState(dim);
        will_free = true;
    } else {
        src = src_out;
    }
    size_t n_qubit = objs.size();
    size_t m_dim = (static_cast<uint64_t>(1) << n_qubit);
    size_t ctrl_mask = 0;
    for (auto& i : ctrls) {
        ctrl_mask |= static_cast<uint64_t>(1) << i;
    }
    std::vector<size_t> obj_masks{};
    for (size_t i = 0; i < m_dim; i++) {
        size_t n = 0;
        size_t mask_j = 0;
        for (size_t j = i; j != 0; j >>= 1) {
            if (j & 1) {
                mask_j += static_cast<uint64_t>(1) << objs[n];
            }
            n += 1;
        }
        obj_masks.push_back(mask_j);
    }
    NQubitsMatrixMul(src, des, obj_masks, ctrl_mask, gate, dim, m_dim);
    if (will_free) {
        derived::FreeState(&src);
    }
}

template <typename derived_, typename calc_type_>
void CPUDensityMatrixPolicyBase<derived_, calc_type_>::NQubitsMatrixMul(const qs_data_p_t& src, const qs_data_p_t& des,
                                                                        std::vector<size_t> obj_masks, size_t ctrl_mask,
                                                                        const matrix_t& gate, index_t dim,
                                                                        size_t m_dim) {
    auto obj_mask = obj_masks.back();
    THRESHOLD_OMP_FOR(
        dim, DimTh, for (omp::idx_t a = 0; a < static_cast<omp::idx_t>(dim); a++) {
            if ((a & obj_mask) == 0) {
                for (index_t b = 0; b <= a; b++) {
                    if ((b & obj_mask) == 0) {
                        bool row_apply = ((a & ctrl_mask) == ctrl_mask);
                        bool col_apply = ((b & ctrl_mask) == ctrl_mask);
                        if (!row_apply && !col_apply) {
                            continue;
                        }
                        VT<VT<qs_data_t>> tmp_mat(m_dim, VT<qs_data_t>(m_dim));
                        if (row_apply) {
                            for (size_t i = 0; i < m_dim; i++) {
                                for (size_t j = 0; j < m_dim; j++) {
                                    for (size_t k = 0; k < m_dim; k++) {
                                        tmp_mat[i][j] += gate[i][k] * GetValue(src, obj_masks[k] | a, obj_masks[j] | b);
                                    }
                                }
                            }
                        } else {
                            for (size_t i = 0; i < m_dim; i++) {
                                for (size_t j = 0; j < m_dim; j++) {
                                    tmp_mat[i][j] = GetValue(src, obj_masks[i] | a, obj_masks[j] | b);
                                }
                            }
                        }
                        if (col_apply) {
                            for (size_t i = 0; i < m_dim; i++) {
                                for (size_t j = 0; j < m_dim; j++) {
                                    qs_data_t new_value = 0;
                                    for (size_t k = 0; k < m_dim; k++) {
                                        new_value += tmp_mat[i][k] * std::conj(gate[j][k]);
                                    }
                                    SetValue(des, obj_masks[i] | a, obj_masks[j] | b, new_value);
                                }
                            }
                        } else {
                            for (size_t i = 0; i < m_dim; i++) {
                                for (size_t j = 0; j < m_dim; j++) {
                                    SetValue(des, obj_masks[i] | a, obj_masks[j] | b, tmp_mat[i][j]);
                                }
                            }
                        }
                    }
                }
            }
        })
}

template <typename derived_, typename calc_type_>
void CPUDensityMatrixPolicyBase<derived_, calc_type_>::ApplyMatrixGate(const qs_data_p_t& src, qs_data_p_t* des_p,
                                                                       const qbits_t& objs, const qbits_t& ctrls,
                                                                       const matrix_t& m, index_t dim) {
    if (objs.size() == 1) {
        derived::ApplySingleQubitMatrix(src, des_p, objs[0], ctrls, m, dim);
    } else if (objs.size() == 2) {
        derived::ApplyTwoQubitsMatrix(src, des_p, objs, ctrls, m, dim);
    } else {
        derived::ApplyNQubitsMatrix(src, des_p, objs, ctrls, m, dim);
    }
}

#ifdef __x86_64__
template struct CPUDensityMatrixPolicyBase<CPUDensityMatrixPolicyAvxFloat, float>;
template struct CPUDensityMatrixPolicyBase<CPUDensityMatrixPolicyAvxDouble, double>;
#elif defined(__amd64)
template struct CPUDensityMatrixPolicyBase<CPUDensityMatrixPolicyArmFloat, float>;
template struct CPUDensityMatrixPolicyBase<CPUDensityMatrixPolicyArmDouble, double>;
#endif
}  // namespace mindquantum::sim::densitymatrix::detail
